# https://www.pythonheidong.com/blog/article/759529/e1a17962c9f524b7f7a9/

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.metrics import pairwise_distances_argmin # 对两个序列中的点进行距离匹配的函数
from sklearn.datasets import load_sample_image # 导入图片数据所用的库
from sklearn.utils import shuffle # 打乱一个有序的序列的函数 # 洗牌

# 实例化导入颐和园图片
china = load_sample_image("china.jpg")
print(len(china[0]))
print(china.shape)

newimage = china.reshape(427*640,3)
print(newimage.shape)


# 该图片中包含了多少种不同的的颜色
df = pd.DataFrame(newimage)
print(df.shape) # (273280, 3) 共有27万多个像素块
df.drop_duplicates(inplace=True)
print(df.shape) # (96615, 3)  只有9万多种颜色

# 可视化图片
# plt.figure(figsize=(15,15))
# plt.imshow(china) # 导入三维数组形成的图片
# plt.show()

# ① 导包

# ② 预处理图像数据，适应SKlearn聚合
# 1.指定簇数
n_cluster = 64
# 2.加载图片
# china = load_sample_image("china.jpg")
# plt.imgshow在浮点数上表现优异，在这里对China数据进行转化，并且归一化处理
# 3.归一化处理图片数据
china = np.array(china,dtype=np.float64)/china.max()
# 4.提取图像格式
w,h,d = original_shape = tuple(china.shape)
# 在这里要确保每个像素块是由3个特征组成的 --- rgb三原色
# assert相当于 raise error if not，表示“不为true就报错”
assert d == 3
# 5.改变图片数据结构以适应KMeans
image_reshape_data = np.reshape(china,(w*h,d))

# ③ KMeans聚类
# 6.数据集过于庞大，先随机抽取1000条进行KMeans聚类，再使用得到的质心去预测整体数据集
sample_data = shuffle(image_reshape_data,random_state=0)[:1000]
k_means = KMeans(n_clusters=n_cluster,random_state=0).fit(sample_data)
# 7.得出质心
centers = k_means.cluster_centers_
print(k_means.cluster_centers_.shape) # (64, 3)
# 8.按照质心预测整体数据集
# labels代表了27万个样本点所对应的簇的质心的索引
labels = k_means.predict(image_reshape_data)
print(labels.shape) # (273280,)

# ④ 质心替换样本，实现矢量量化
# 9.使用质心替换所有的样本
image_kmeans = image_reshape_data.copy()
# 10.遍历替换质心数
for i in range(w*h):
    image_kmeans[i] = centers[labels[i]]
# 检验是否替换成功
print(pd.DataFrame(image_kmeans).drop_duplicates().shape) # (64, 3)

# 11.查看矢量量化后的图像
# 11.1 恢复图片结构
image_kmeans = image_kmeans.reshape(w,h,d)
print(image_kmeans.shape) # (427, 640, 3)

# ⑤ 随机的矢量量化
'''
随机的矢量量化
'''
# 1.产生随机质心
centroid_random = shuffle(image_reshape_data, random_state=0)[:n_cluster]
print(centroid_random)
# [[0.92156863 0.9254902  0.94509804]
#  [0.37647059 0.37647059 0.14117647]
#   .........
#  [0.40784314 0.48235294 0.12941176]
#  [0.81568627 0.86666667 0.8       ]]

# 函数pairwise_distances_argmin(x1,x2,axis)
# 2.用来计算x2中的每个样本到x1中的每个样本点的距离,并返回和x2相同形状的,x1中对应的最近的样本点的对应聚类数
labels_random = pairwise_distances_argmin(centroid_random, image_reshape_data, axis=0)
print(labels_random)  # [55 55 55 ... 52 60 60]
print(labels_random.shape) # (273280,)
print(len(set(labels_random))) # 64

# 3.使用随机质心来替换所有样本
image_random = image_reshape_data. copy()
print(image_random)
# [[0.68235294 0.78823529 0.90588235]
#  [0.68235294 0.78823529 0.90588235]
#  [0.68235294 0.78823529 0.90588235]
#  ...
#  [0.16862745 0.19215686 0.15294118]
#  [0.05098039 0.08235294 0.02352941]
#  [0.05882353 0.09411765 0.02745098]]
for i in range(w*h):
   image_random[i] = centroid_random [labels_random[i]]
# 4.恢复图片的结构
image_random= image_random.reshape(w,h,d)
print(image_random.shape)

# ⑥ 可视化
plt. figure(figsize=(10,10))
plt. axis('off')
plt. title('Original image (96,615 colors)')
plt. imshow(china)

plt. figure(figsize=(10,10))
plt. axis('off')
plt. title('Quantized image (64 colors, K-Means)')
plt. imshow(image_kmeans)

plt. figure(figsize=(10,10))
plt. axis('off')
plt. title('Quantized image (64 colors, Random)')
plt. imshow(image_random)

plt. show()
